//! block-chain DataBase Storage
//!
//! This module provides a simple interface for interacting with a RocksDB database.
//! It allows you to store and retrieve key-value pairs, as well as perform other operations
//! such as checking if a key exists or deleting a key.
//!

use std::{path::Path, sync::Arc};

use tokio::task;

pub type DBType = rocksdb::DB;
// pub type BcDBType = Arc<RwLock<DBType>>;
pub type BcDBType = Arc<DBType>; // rocksdb::DB -- muilti-thread safe

pub use rocksdb::IteratorMode as DBIteratorMode;
pub use rocksdb::Options as DBOptions;

#[derive(thiserror::Error, Debug)]
pub enum DBError {
    #[error("An error occurred while interacting with the database: {0}")]
    DataBaseError(#[from] rocksdb::Error),

    #[error("Join error: {0}")]
    JoinError(#[from] tokio::task::JoinError),

    #[error("IO error: {0}")]
    IOError(#[from] std::io::Error),
}

impl PartialEq for DBError {
    fn eq(&self, other: &Self) -> bool {
        matches!(
            (self, other),
            (DBError::DataBaseError(_), DBError::DataBaseError(_))
                | (DBError::JoinError(_), DBError::JoinError(_))
                | (DBError::IOError(_), DBError::IOError(_))
        )
    }
}

#[derive(Debug)]
pub struct BaseBcDb {
    pub db: BcDBType,
    path: std::path::PathBuf,
}

impl BaseBcDb {
    pub fn new(path: &Path) -> Result<Self, DBError> {
        let db = crate::open_db(path)?;
        let db = Arc::new(db);

        Ok(Self {
            db,
            path: path.to_path_buf(),
        })
    }

    pub fn destroy_db_in_disk(path: &Path) {
        destroy_db(path).unwrap();
    }

    pub fn get_db_path(&self) -> &Path {
        &self.path
    }

    pub async fn put(&self, key: &str, value: Vec<u8>) -> Result<(), DBError> {
        let db = self.db.clone();
        let key = key.to_string();
        if let Err(e) = task::spawn_blocking(move || db.put(key, value)).await? {
            return Err(DBError::DataBaseError(e));
        }
        Ok(())
    }
    pub async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, DBError> {
        let db = self.db.clone();
        let key = key.to_string();
        let val = task::spawn_blocking(move || db.get(key)).await?;

        match val {
            Ok(Some(v)) => Ok(Some(v.to_vec())),
            Ok(None) => Ok(None),
            Err(e) => Err(DBError::DataBaseError(e)),
        }
    }
    pub async fn erase(&self, key: &str) -> Result<(), DBError> {
        let db = self.db.clone();
        let key = key.to_string();
        if let Err(e) = task::spawn_blocking(move || db.delete(key)).await? {
            return Err(DBError::DataBaseError(e));
        }
        Ok(())
    }

    pub async fn exists(&self, key: &str) -> bool {
        let db = self.db.clone();
        let key = key.to_string();
        let val = task::spawn_blocking(move || {
            // db.key_may_exist(key) // MARK
            db.get(key).unwrap().is_some()
        })
        .await;

        match val {
            Ok(v) => v,
            Err(e) => {
                tracing::error!("Error while checking if key exists: {:?}", e);
                false
            }
        }
    }

    // Rust 的 rocksdb 库为 DB 类型实现了 Drop trait，当数据库实例离开作用域时，会自动关闭并释放资源。
    // 这里只是 flush db
    pub async fn flush(&self) -> Result<(), DBError> {
        let db = self.db.clone();
        if let Err(e) = task::spawn_blocking(move || db.flush()).await? {
            return Err(DBError::DataBaseError(e));
        }
        Ok(())
    }

    pub async fn get_item_number(&self) -> Result<usize, DBError> {
        let db = self.db.clone();
        let val = task::spawn_blocking(move || {
            let iter = db.iterator(DBIteratorMode::Start);
            let mut count = 0;
            for _ in iter {
                count += 1;
            }
            count
        })
        .await?;

        Ok(val)
    }
}

fn open_db(path: &Path) -> Result<DBType, DBError> {
    let mut opts = DBOptions::default();
    opts.create_if_missing(true);
    let db = DBType::open(&opts, path)?;
    Ok(db)
}

fn destroy_db(path: &Path) -> Result<(), DBError> {
    let opts = DBOptions::default();
    DBType::destroy(&opts, path)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    use scopeguard::defer;

    #[test]
    fn test_open_db() {
        let tmp_dir = bc_util::convert_to_temp_dir("test_open_db");
        let path = tmp_dir.path().to_path_buf();
        defer! {
            tmp_dir.close().unwrap();
            assert_eq!(path.exists(), false);
        };

        let db = open_db(&path).unwrap();
        assert!(!db.key_may_exist("test_key"));
        assert!(path.exists());
    }

    async fn test_iter(db: &BaseBcDb, key1: &str, value1: &Vec<u8>, key2: &str, value2: &Vec<u8>) {
        for (index, result) in db.db.iterator(DBIteratorMode::Start).enumerate() {
            let (k, v) = result.unwrap();
            println!("key: {:?}, value: {:?}", k, v);
            let k = std::str::from_utf8(&k).unwrap();
            let v = v.into_vec();

            if index == 0 {
                assert_eq!(k, key1);
                assert_eq!(v, *value1);
            } else if index == 1 {
                assert_eq!(k, key2);
                assert_eq!(v, *value2);
            } else {
                panic!("Unexpected key");
            }
        }
    }

    #[tokio::test]
    async fn test_base_bc_db() {
        let key1 = "test_key1";
        let value1: Vec<u8> = vec![1, 2, 3];
        let key2 = "test_key2";
        let value2: Vec<u8> = vec![4, 5, 6];

        let value1_1: Vec<u8> = vec![7, 8, 9];

        let tmp_dir = bc_util::convert_to_temp_dir("test_base_bc_db");
        let path = tmp_dir.path().to_path_buf();

        {
            let db = BaseBcDb::new(&path).unwrap();
            assert_eq!(db.get_db_path(), &path);

            assert_eq!(db.get_item_number().await.unwrap(), 0);

            assert!(!(db.exists(key1).await));
            db.put(key1, value1.clone()).await.unwrap(); // insert
            assert!(db.exists(key1).await);
            assert_eq!(db.get(key1).await.unwrap(), Some(value1.clone()));
            db.put(key1, value1_1.clone()).await.unwrap(); // update
            assert_eq!(db.get(key1).await.unwrap(), Some(value1_1.clone()));

            db.put(key2, value2.clone()).await.unwrap(); // insert
            assert_eq!(db.get(key2).await.unwrap(), Some(value2.clone()));
            assert!(db.exists(key2).await);

            assert_eq!(db.get_item_number().await.unwrap(), 2);

            test_iter(&db, key1, &value1_1, key2, &value2).await; // iter

            db.erase(key2).await.unwrap(); // delete

            assert_eq!(db.get_item_number().await.unwrap(), 1);
            assert!(!(db.exists(key2).await));
        }

        BaseBcDb::destroy_db_in_disk(&path);
        assert!(!path.exists());
    }

    // new, put, flush, exist
    #[tokio::test]
    async fn test_bc_db_put_flush_exist() {
        let tmp_dir = bc_util::convert_to_temp_dir("test_bc_db_put_flush_exist");
        let path = tmp_dir.path().to_path_buf();
        defer! {
            tmp_dir.close().unwrap();
        };

        let key1 = "test_key1";
        let value1: Vec<u8> = vec![1, 2, 3];
        let key2 = "test_key2";

        let db = BaseBcDb::new(&path).unwrap();

        assert!(db.put(key1, value1.clone()).await.is_ok());
        assert!(!db.exists(key2).await);
        assert!(db.flush().await.is_ok());
        assert!(!db.exists(key2).await);
    }
}
